#ifndef REGALLOC_H
#define REGALLOC_H

#include "flow.h"
#include "opcode.h"
#include "ssa.h"
#include "lir.h"
#include <set>
#include <vector>
#include <list>
#include <deque>
#include <map>
#include <iterator>
#include <algorithm>
#include <cassert>

/// the wrapper for RegOperand in lir
/// make it convenient for regalloc
class RegAllocOperand
{
private:
    std::set<int> adjList;

    int id;

    /// the id number of next virtual operand
    static int nextId;

    /// disable creating RegAllocOperand without register info
    RegAllocOperand();
public:
    int degree;
    RegAllocOperand* alias;

    // the machine register assigned for this operand
    int realRegister;

    /// the associated move lirs
    std::set<Lir*> moveList;       // TODO: refactor, encapsulate this one

    /// RegAllocOperand is also identified by operand's id
    RegAllocOperand(int num):realRegister(-1)
    {
        id = num;
        degree = 0;
        alias = 0;
    }

    bool interfere(int num)
    {
        return adjList.find(num) != adjList.end();
    }

    bool interfere(RegAllocOperand *operand)
    {
        return adjList.find(operand->getId()) != adjList.end();
    }

    void addAdj(int num)
    {
        adjList.insert(num);
        degree++;
    }

    void removeAdj(int num)
    {
        adjList.erase(num);
        degree--;
    }

    void addMove(Lir *movLir)
    {
        assert(movLir->type == MOV && "add non-mov lir to moveList");
        moveList.insert(movLir);
    }

    void removeMove(Lir *movLir)
    {
        moveList.erase(movLir);
    }


    /// return (moveList ^  (activeMoves   V   worklistMoves))
    std::set<Lir*> nodeMoves(std::set<Lir*>& activeAndWorklist)
    {
        std::set<Lir*> tmpSet;

        set_intersection(moveList.begin(),
                         moveList.end(),
                         activeAndWorklist.begin(),
                         activeAndWorklist.end(),
                         inserter(tmpSet,tmpSet.begin()));

        return tmpSet;
    }

    bool isMoveRelated(std::set<Lir*>& activeAndWorklist)
    {
        return nodeMoves(activeAndWorklist).size() != 0;
    }

    int getId()
    {
        return id;
    }

    std::set<int> getAdjlist()
    {
        // TODO: should be adjList ^ (selectStack V colescedNodes)
        return std::set<int>(adjList);
    }

    /// set nextId that used for next made operand
    static void setTotalReg(int n)
    {
        nextId = n + 1;
    }

    static RegAllocOperand* make()
    {
        return new RegAllocOperand(nextId++);
    }

    /// does the lir reference this operand
    bool referencedBy(Lir *lir)
    {
        RegOperand *regDst = dynamic_cast<RegOperand*>(lir->dst);
        RegOperand *regArg1 = dynamic_cast<RegOperand*>(lir->arg1);
        RegOperand *regArg2 = dynamic_cast<RegOperand*>(lir->arg2);

        if((regDst && id == *regDst) ||
           (regArg1 && id == *regArg1) ||
           (regArg2 && id == *regArg2)){
            return true;
        }

        return false;
    }

};

/// the node contains liveness info for a basicblock
class RegAllocNode : public AnalysisNode
{
private:
    // currently only x86 is supported
    static X86Encoder encoder;

public:
    // TODO: refactor it to liveness ananlysis
    dyn_bitset *liveIn;
    dyn_bitset *liveOut;
    dyn_bitset *gen;
    dyn_bitset *killed;

    LirQueue lirs;
    typedef LirQueue::iterator lir_iterator;
    typedef LirQueue::reverse_iterator lir_reverse_iterator;

    RegAllocNode(BasicBlock *b,int regNum)
        :AnalysisNode(b)
    {
        liveIn = liveOut = gen = killed = 0;
        reset(regNum);

        // transform opcodes to lirs
        encoder.encode(b);
        X86Encoder::fixupLabels();

        lirs = encoder.getLirs();
    }

    void reset(int regNum)
    {
        if(liveIn){
            delete liveIn;
        }
        if(liveOut){
            delete liveOut;
        }
        if(gen){
            delete gen;
        }
        if(killed){
            delete killed;
        }

        liveIn = new dyn_bitset(regNum,false);
        liveOut = new dyn_bitset(regNum,false);
        gen = new dyn_bitset(regNum,false);
        killed = new dyn_bitset(regNum,false);
    }

    virtual ~RegAllocNode()
    {
        delete liveIn;
        delete liveOut;
        delete gen;
        delete killed;
    }

    void define(int regIndex)
    {
        if(regIndex < 0){
            return;
        }

        if(!(*gen)[regIndex]){
            (*killed)[regIndex] = true;
        }
    }

    void use(int regIndex)
    {
        if(regIndex < 0){
            return;
        }

        if(!(*killed)[regIndex]){
            (*gen)[regIndex] = true;
        }
    }

};

/// the main class for Register Allocation
class RegAlloc : public Analysis<RegAllocNode>
{
private:
    std::map<int,RegAllocOperand*> initial;
    std::list<RegAllocOperand*> simplifyWorklist;
    std::list<RegAllocOperand*> freezeWorklist;
    std::list<RegAllocOperand*> spillWorklist;
    std::list<RegAllocOperand*> spilledNodes;
    std::list<RegAllocOperand*> coalescedNodes;
    std::list<RegAllocOperand*> coloredNodes;
    std::list<RegAllocOperand*> selectStack;

    std::set<Lir*> coalescedMoves;
    std::set<Lir*> constrainedMoves;
    std::set<Lir*> frozenMoves;
    std::set<Lir*> worklistMoves;
    std::set<Lir*> activeMoves;

    static int totalRegs;
    static int getTotalRegs()
    {
        return totalRegs;
    }

    int K;

    /// refactor: should be arch-dependent
    unsigned int availRegMask;

    // replace lir's dst register(spilled) with tempId
    static bool replaceDst(Lir* lir,int spilledId,int tempId);

    // replace lir's args register(spilled) with tempId
    static bool replaceArgs(Lir* lir,int spilledId,int tempId);
    int getDefs(Lir *lir);
    int getUses(Lir *lir,int uses[]);

    void scanBlocks();
    void addEdge(dyn_bitset& adjs,int regNum);
    void addEdge(std::set<int>& adjs,int regNum);
    void addWorklist(RegAllocOperand* operand);
    void enableMoves(std::set<int>& regs);
    void decrementDegree(RegAllocOperand *operand);
    void combine(RegAllocOperand *u,RegAllocOperand *v);
    RegAllocOperand* getAlias(RegAllocOperand *operand);
    void freezeMoves(RegAllocOperand *operand);
    bool isSafeCoalesce(RegAllocOperand *v,RegAllocOperand *u);

    void initialize();
    void addOperand(int regNum);
    void collectOperands();
    void build();
    void makeWorklist();
    void simplify();
    void coalesce();
    void freeze();
    void selectSpill();
    void assignColors();
    void rewriteProgram();
    void removeRedundantMoves();
    void assignRealRegister();

    template<class T>
    T setUnion(const T& set1,const T& set2)
    {
        T tmpSet;

        set_union(set1.begin(),
                  set1.end(),
                  set2.begin(),
                  set2.end(),
                  inserter(tmpSet,tmpSet.begin()));
        return tmpSet;
    }

public:
    RegAlloc(CFG *c)
        :Analysis<RegAllocNode>(c)
    {
        log.setDebug(false);
        K = 3;
        availRegMask = (1 << K) - 1;
        c->regalloc = this;
    }

    virtual RegAllocNode* createNode(BasicBlock *bb);

    // TODO: add a proxy for this interface
    LirQueue& getLirsFor(BasicBlock *bb);


    void analyzeLiveOut();
    void allocateRegister();
    void printBlocks();
    void printOperands();
    void printLirs();
    void printAllocatedCode();
};

#endif

